from nl2sql.main import data_info
from nl2sql.model.llm import LLM
from typing import Optional, List
from nl2sql.prompt.query_analysis_prompt import sql_analysis_prompt
from nl2sql.tools.extract_block import extract_python


class QueryAnalysis:
    def __init__(self,
                 model: LLM = None) -> None:

        self.prompt = sql_analysis_prompt
        # 给这个类赋予一个属性名为llm，这个属性的数据类型是LLM，或者是None
        self.llm: Optional[LLM] = model

    def analyze(self, query: str, data_info: list[str]) -> Optional[str]:
        """
        用于拆解用户的问题，为子问题
        :param query:
        :return: None or List of String
        """
        # 把Prompt模板，进行填充
        filled_prompt = self.prompt.format(query=query,data_info=data_info)
        # 让大模型对填充好的模板，进行推理
        llm_resp: str = self.llm.chat(prompt=filled_prompt)
        # 对大模型的输出，进行后处理（因为我们没办法保证大模型输出的格式就是我们想要的）
        result = llm_resp
        return result

